#!/usr/bin/python3
# -*- coding: UTF-8 -*-
# __author__ = 'zd'
import sys


def get_ner_measure(golden_lists, predict_lists, label_type='BMES'):
    """
    input as sentence level labels
    :param golden_lists: [sentence_number, sentence_length] -> element 'B-PER'
    :param predict_lists: [sentence_number, sentence_length] -> element 'B-PER'
    :param label_type: label schema -> choice 'BMES' or 'BIO'
    :return: accuracy precision recall f_score
    """
    sentence_number = len(golden_lists)
    golden_full = []
    predict_full = []
    right_full = []
    right_tag = 0
    all_tag = 0
    for idx in range(sentence_number):
        # word_list = sentence_lists[idx]
        golden_list = golden_lists[idx]
        predict_list = predict_lists[idx]
        for idy in range(len(golden_list)):
            if golden_list[idy] == predict_list[idy]:
                right_tag += 1
        all_tag += len(golden_list)
        if label_type == 'BMES':
            gold_matrix = get_ner_BMES(golden_list)
            pred_matrix = get_ner_BMES(predict_list)
        else:
            gold_matrix = get_ner_BIO(golden_list)
            pred_matrix = get_ner_BIO(predict_list)
        right_ner = list(set(gold_matrix).intersection(set(pred_matrix)))
        golden_full += gold_matrix
        predict_full += pred_matrix
        right_full += right_ner
    right_num = len(right_full)
    golden_num = len(golden_full)
    predict_num = len(predict_full)
    if predict_num == 0:
        precision = -1
    else:
        precision = (right_num + 0.0) / predict_num
    if golden_num == 0:
        recall = -1
    else:
        recall = (right_num+0.0) / golden_num
    if precision == -1 or recall == -1 or precision + recall <= 0.:
        f_score = -1
    else:
        f_score = 2 * precision * recall / (precision+recall)
    accuracy = (right_tag+0.0) / all_tag
    # print('Accuracy: ', right_tag, '/', all_tag, '=', accuracy)
    # print('gold_num = ', golden_num, ' pred_num = ', predict_num, ' right_num = ', right_num)
    return accuracy, precision, recall, f_score


def reverse_style(input_string):
    target_position = input_string.index('[')
    input_len = len(input_string)
    output_string = input_string[target_position: input_len] + input_string[0: target_position]
    return output_string


def get_ner_BMES(label_list):
    # list_len = len(word_list)
    # assert(list_len == len(label_list)), 'word list size not match with label list'
    list_len = len(label_list)
    begin_label = 'B-'
    end_label = 'E-'
    single_label = 'S-'
    whole_tag = ''
    index_tag = ''
    tag_list = []
    stand_matrix = []
    for i in range(list_len):
        # word_label = word_list[i]
        current_label = label_list[i].upper()
        if begin_label in current_label:
            if index_tag != '':
                tag_list.append(whole_tag + ',' + str(i - 1))
            whole_tag = current_label.replace(begin_label, '', 1) + '[' + str(i)
            index_tag = current_label.replace(begin_label, '', 1)
        elif single_label in current_label:
            if index_tag != '':
                tag_list.append(whole_tag + ',' + str(i - 1))
            whole_tag = current_label.replace(single_label, '', 1) + '[' + str(i)
            tag_list.append(whole_tag)
            whole_tag = ''
            index_tag = ''
        elif end_label in current_label:
            if index_tag != '':
                tag_list.append(whole_tag + ',' + str(i))
            whole_tag = ''
            index_tag = ''
        else:
            continue
    if (whole_tag != '') & (index_tag != ''):
        tag_list.append(whole_tag)
    tag_list_len = len(tag_list)

    for i in range(tag_list_len):
        if len(tag_list[i]) > 0:
            tag_list[i] = tag_list[i] + ']'
            insert_list = reverse_style(tag_list[i])
            stand_matrix.append(insert_list)
    # print(stand_matrix)
    return stand_matrix


def get_ner_BIO(label_list):
    # list_len = len(word_list)
    # assert(list_len == len(label_list)), 'word list size not match with label list'
    list_len = len(label_list)
    begin_label = 'B-'
    inside_label = 'I-' 
    whole_tag = ''
    index_tag = ''
    tag_list = []
    stand_matrix = []
    for i in range(list_len):
        # word_label = word_list[i]
        current_label = label_list[i].upper()
        if begin_label in current_label:
            if index_tag == '':
                whole_tag = current_label.replace(begin_label, '', 1) + '[' + str(i)
                index_tag = current_label.replace(begin_label, '', 1)
            else:
                tag_list.append(whole_tag + ',' + str(i - 1))
                whole_tag = current_label.replace(begin_label, '', 1) + '[' + str(i)
                index_tag = current_label.replace(begin_label, '', 1)
        elif inside_label in current_label:
            if current_label.replace(inside_label, '', 1) == index_tag:
                whole_tag = whole_tag 
            else:
                if (whole_tag != '') & (index_tag != ''):
                    tag_list.append(whole_tag + ',' + str(i - 1))
                whole_tag = ''
                index_tag = ''
        else:
            if (whole_tag != '') & (index_tag != ''):
                tag_list.append(whole_tag + ',' + str(i - 1))
            whole_tag = ''
            index_tag = ''

    if (whole_tag != '') & (index_tag != ''):
        tag_list.append(whole_tag)
    tag_list_len = len(tag_list)

    for i in range(tag_list_len):
        if len(tag_list[i]) > 0:
            tag_list[i] = tag_list[i] + ']'
            insert_list = reverse_style(tag_list[i])
            stand_matrix.append(insert_list)
    return stand_matrix


def readSentence(input_file):
    with open(input_file, 'r', encoding='UTF-8') as f:
        sentences = []
        labels = []
        sentence = []
        label = []
        for line in f:
            if len(line) < 2:
                sentences.append(sentence)
                labels.append(label)
                sentence = []
                label = []
            else:
                pair = line.strip('\n').split(' ')
                sentence.append(pair[0])
                label.append(pair[-1])
        return sentences, labels


def readTwoLabelSentence(input_file, pred_col=-1):
    with open(input_file, 'r', encoding='UTF-8') as f:
        sentences = []
        predict_labels = []
        golden_labels = []
        sentence = []
        predict_label = []
        golden_label = []
        for line in f:
            if '##score##' in line:
                continue
            if len(line) < 2:
                sentences.append(sentence)
                golden_labels.append(golden_label)
                predict_labels.append(predict_label)
                sentence = []
                golden_label = []
                predict_label = []
            else:
                pair = line.strip('\n').split(' ')
                sentence.append(pair[0])
                golden_label.append(pair[1])
                predict_label.append(pair[pred_col])

        return sentences, golden_labels, predict_labels


def measure_from_file(golden_file, predict_file, label_type='BMES'):
    print('Get measure from file:', golden_file, predict_file)
    print('Label format:', label_type)
    golden_sent, golden_labels = readSentence(golden_file)
    predict_sent, predict_labels = readSentence(predict_file)
    accuracy, precision, recall, f_score = get_ner_measure(golden_labels, predict_labels, label_type)
    print("accuracy: %s, precision: %s recall: %s, f_score: %s" % (accuracy, precision, recall, f_score))


def fmeasure_from_singlefile(two_label_file, label_type='BMES', pred_col=-1):
    sent,golden_labels,predict_labels = readTwoLabelSentence(two_label_file, pred_col)
    accuracy, precision, recall, f_score = get_ner_measure(golden_labels, predict_labels, label_type)
    print("accuracy: %s, precision: %s recall: %s, f_score: %s" % (accuracy, precision, recall, f_score))
